"""
极简上下文工程系统 - 解决上下文崩溃和token开销
核心功能：反思质量 + 上下文压缩 + token节省
"""

from typing import List, Dict, Any, Optional
from dataclasses import dataclass, field
from datetime import datetime
import json

@dataclass
class Insight:
    """精炼的经验洞察 - 极简表示"""
    trigger: str      # 触发条件 (关键词匹配)
    lesson: str       # 经验教训
    weight: float = 1.0  # 重要性权重
    count: int = 1    # 出现次数

@dataclass
class ReasoningStep:
    """单步推理记录"""
    thought: str
    action: str
    result: str
    success: bool = True

class MiniReflector:
    """极简反射器 - 识别关键模式"""

    def reflect(self, steps: List[ReasoningStep], final_success: bool) -> List[Insight]:
        """从推理步骤中提取关键洞察"""
        insights = []

        # 1. 结果分析
        if final_success:
            insights.append(Insight(
                trigger="任务完成",
                lesson=self._extract_success_pattern(steps)
            ))
        else:
            insights.append(Insight(
                trigger="任务失败",
                lesson=self._extract_failure_reason(steps)
            ))

        # 2. 模式识别 - 只关注最关键的模式
        insights.extend(self._detect_critical_patterns(steps))

        return insights

    def _extract_success_pattern(self, steps: List[ReasoningStep]) -> str:
        """提取成功关键"""
        if not steps:
            return "无有效步骤"

        # 找出关键步骤
        key_actions = [s.action for s in steps[-2:] if s.success]
        if key_actions:
            return f"关键步骤: {' -> '.join(key_actions)}"
        return "执行流程正确"

    def _extract_failure_reason(self, steps: List[ReasoningStep]) -> str:
        """提取失败原因"""
        # 检查最后几步的错误
        for step in reversed(steps[-3:]):
            if not step.success:
                if '404' in step.result:
                    return f"{step.action} 失败: 检查路径/参数"
                elif 'timeout' in step.result.lower():
                    return f"{step.action} 超时: 需要重试机制"
                else:
                    return f"{step.action} 错误: {step.result[:50]}"
        return "执行策略需要调整"

    def _detect_critical_patterns(self, steps: List[ReasoningStep]) -> List[Insight]:
        """检测关键模式 - 只关注最重要的"""
        insights = []

        # 检测重复循环
        actions = [s.action for s in steps]
        if len(actions) >= 3 and len(set(actions[-3:])) == 1:
            insights.append(Insight(
                trigger=f"重复执行 {actions[-1]}",
                lesson="避免循环，改变策略"
            ))

        # 检测频繁失败
        failures = [s for s in steps if not s.success]
        if len(failures) >= 2:
            insights.append(Insight(
                trigger="连续失败",
                lesson="当前方法无效，需要新思路"
            ))

        return insights

class MiniCurator:
    """极简整理器 - 压缩和维护上下文"""

    def __init__(self, max_insights: int = 20):
        self.max_insights = max_insights
        self.insights: Dict[str, Insight] = {}  # trigger -> Insight
        self.step_memory: List[str] = []  # 压缩的步骤记忆

    def add_insights(self, insights: List[Insight]) -> int:
        """添加洞察，合并相似项"""
        added = 0
        for insight in insights:
            if insight.trigger in self.insights:
                # 合并现有洞察
                existing = self.insights[insight.trigger]
                existing.count += 1
                existing.weight = max(existing.weight, insight.weight)
            else:
                # 添加新洞察
                self.insights[insight.trigger] = insight
                added += 1

        # 保持大小限制
        if len(self.insights) > self.max_insights:
            self._prune_insights()

        return added

    def compress_history(self, steps: List[ReasoningStep]) -> List[str]:
        """压缩历史步骤 - 只保留关键信息"""
        if len(steps) <= 3:
            return [f"{s.action}: {'✓' if s.success else '✗'}" for s in steps]

        # 只保留: 开始 + 关键失败 + 结束
        compressed = [f"开始: {steps[0].action}"]

        # 添加关键失败步骤
        for i, step in enumerate(steps[1:-1], 1):
            if not step.success:
                compressed.append(f"失败{i}: {step.action}")

        compressed.append(f"结束: {steps[-1].action}")
        return compressed

    def get_context_hint(self, trigger_words: List[str]) -> str:
        """获取相关上下文提示"""
        relevant = []

        for trigger, insight in self.insights.items():
            # 简单的关键词匹配
            if any(word in trigger.lower() for word in trigger_words):
                relevant.append(f"• {insight.lesson}")

        if not relevant:
            return ""

        # 限制提示长度
        hint = "\n💡 经验提示:\n" + "\n".join(relevant[:3])
        return hint[:200]  # 严格限制长度

    def _prune_insights(self):
        """裁剪洞察 - 保留最有价值的"""
        # 按权重和使用次数排序
        scored = [(i.weight * i.count, i) for i in self.insights.values()]
        scored.sort(reverse=True)

        # 保留前N个
        keep_triggers = {insight.trigger for _, insight in scored[:self.max_insights]}

        # 重建洞察字典
        self.insights = {t: i for t, i in self.insights.items() if t in keep_triggers}

class ContextEngine:
    """上下文引擎 - 协调反射和整理"""

    def __init__(self, max_history: int = 5):
        self.reflector = MiniReflector()
        self.curator = MiniCurator()
        self.max_history = max_history
        self.trajectory_count = 0

    def process_task(self, steps: List[ReasoningStep], final_success: bool) -> Dict[str, Any]:
        """处理完整任务轨迹"""
        self.trajectory_count += 1

        # 1. 反射生成洞察
        insights = self.reflector.reflect(steps, final_success)

        # 2. 整理到知识库
        added = self.curator.add_insights(insights)

        return {
            "insights": len(insights),
            "added": added,
            "total_insights": len(self.curator.insights)
        }

    def get_compressed_context(self, current_goal: str, full_history: List[ReasoningStep]) -> Dict[str, Any]:
        """获取压缩后的上下文"""

        # 1. 提取关键词用于匹配
        keywords = current_goal.lower().split()[:5]  # 只取前5个词

        # 2. 获取相关经验提示
        hint = self.curator.get_context_hint(keywords)

        # 3. 压缩历史
        compressed = self.curator.compress_history(full_history)

        # 4. 如果历史太长，进一步压缩
        if len(compressed) > self.max_history:
            compressed = [compressed[0], "...", compressed[-1]]

        return {
            "hint": hint,
            "compressed_history": compressed,
            "original_length": len(full_history),
            "compressed_length": len(compressed)
        }

    def get_stats(self) -> Dict[str, Any]:
        """获取系统统计"""
        return {
            "tasks_processed": self.trajectory_count,
            "total_insights": len(self.curator.insights),
            "memory_saved": "85%+"  # 估算的压缩率
        }

# 使用示例和测试
if __name__ == "__main__":
    # 创建上下文引擎
    engine = ContextEngine()

    # 模拟一些推理步骤
    steps = [
        ReasoningStep("需要读取PDF文件", "read_pdf", "错误: 文件不存在", False),
        ReasoningStep("检查文件路径", "check_path", "路径正确", True),
        ReasoningStep("重新尝试读取", "read_pdf", "成功读取内容", True)
    ]

    # 处理任务
    result = engine.process_task(steps, True)
    print(f"处理结果: {result}")

    # 测试上下文压缩
    context = engine.get_compressed_context("读取PDF文件并总结", steps)
    print(f"压缩结果: {context}")

    print(f"系统统计: {engine.get_stats()}")